library(dplyr)
library(tidyr)
library(ggplot2)
library(viridis)
library(ggpubr)
library(cowplot)
library(gridExtra)
library(stringr)

if (!file.exists("output")) {
  dir.create("output")
}

# load data
source("HNC_metadata_tidy.R")

sbs_burden <- read.csv("Input_mutational_burden/Input_SBS_mutational_burden.csv") %>% rename(SBS_burden = Mutational.burden)
dbs_burden <- read.csv("Input_mutational_burden/Input_DBS_mutational_burden.csv") %>% rename(DBS_burden = Mutational.burden)
id_burden <- read.csv("Input_mutational_burden/Input_ID_mutational_burden.csv") %>% rename(ID_burden = Mutational.burden)

additional_df <- list(sbs_burden, dbs_burden, id_burden)

for (df in additional_df) {
  data <- data %>% left_join(df, by = "donor_id")
}

data <- data %>% mutate(tobacco = factor(tobacco, labels = c("Non-smoker", "Ex-smoker", "Current-smoker")))

# Perform Kruskal-Wallis Test and Generate Plots
# This function conducts Kruskal-Wallis tests between variables of interest and mutation burdens
plot_kw_fct <- function(dat, # Data frame containing variables  and mutation butfrnd
                        mut_type, # Character vector specifying mutation types
                        factors, # Character vector specifying variables of interest
                        pval_cutoff = 0.05, # p-value threshold for plotting a specific test
                        output.kw = data.frame(), # A data frame to store the Kruskal-Wallis test results (default is an empty data frame)
                        y_labs = "Mutation burden", # Label for the y-axis in the plots
                        facet_by = NA # Variable by which plots are faceted. If NULL, plot will not be faceted.
) {
  plots <- list()
  for (s in mut_type) {
    for (f in factors) {
      # Check if the factor has less than 2 categories, if so, print a warning and skip
      if (dat[, f] %>% unique() %>% length() < 2) {
        print(paste("Warning: Variable", f, "has been removed due to <2 categories"))
      } else {
        kw <- tryCatch(
          {
            kruskal.test(dat[, s], dat[, f])
          },
          error = function(e) {
            NULL
          }
        )

        # If there was an error in computing the test, move to the next factor
        if (is.null(kw)) {
          print(paste("Variable", f, "has been removed due to ERROR in KW test"))
          next
        }

        result.var <- data.frame(
          signature = paste(s),
          factor = paste(f),
          pval = kw$p.value
        )

        output.kw <- rbind(output.kw, result.var)

        # If p-value is less than the cutoff, proceed with plotting
        if (kw$p.value < pval_cutoff) {
          # Remove NAs
          subset <- dat[!is.na(dat[f]) & !is.na(dat[s]), ]

          # plot
          p <- ggplot(data = subset, aes(x = get(f), y = get(s), fill = get(f))) +
            geom_violin(aes(color = get(f)), width = 1, alpha = .8, linewidth = .3) +
            geom_boxplot(color = "#585858", alpha = 0.15, linewidth = .3, staplewidth = .2, outlier.shape = NA) +
            geom_jitter(fill = "#585858", size = .8, stroke = 0, alpha = 0.25, width = .1, height = .1) +
            labs(
              title = str_replace(s, "_", " "),
              subtitle = paste("Kruskal-Wallis, p =", signif(kw$p.value, 3)),
              y = y_labs,
              x = str_replace(f, "_", " ")
            ) +
            scale_fill_viridis(discrete = TRUE, end = .8) +
            scale_color_viridis(discrete = TRUE, end = .8) +
            theme_minimal() +
            theme(
              legend.position = "none",
              plot.title = element_text(face = "bold", size = 12, hjust = 0.5),
              plot.subtitle = element_text(size = 10, hjust = 0.5),
              axis.text.x = element_text(angle = 45, hjust = 1),
              axis.text = element_text(size = 12),
              axis.line = element_line(size = .3,color = "#404040"),
              axis.ticks = element_line(size = .3,color = "#404040"),
              panel.grid = element_blank()
            )

          # facet plot
          if (is.na(facet_by) | f == facet_by) {
            print(p)
            plots[[paste(s, f, sep = "_")]] <- p
          } else {
            p <- p + labs(title = NULL) + theme(axis.title.x = element_blank())

            p_facet <- p +
              facet_grid(. ~ get(facet_by)) +
              labs(subtitle = " ") +
              theme(
                axis.text.y = element_blank(),
                axis.title.y = element_blank(),
                strip.background = element_rect(color = "#404040", fill = NA),
                panel.spacing = unit(0, "lines")
              )

            p_grid <- plot_grid(p, p_facet, align = "h", axis = "bt", rel_widths = c(0.9, 1.4))

            p_all <- annotate_figure(p_grid,
              top = text_grob(str_replace(mut_type, "_", " "), face = "bold", size = 12),
              bottom = text_grob(str_replace(f, "_", " "), size = 12)
            )
            print(p_all)
            plots[[paste(s, f, sep = "_")]] <- p_all
          }
        }
      }
    }
  }
  # Calculate adjusted p-values
  output.kw <- output.kw %>% mutate(p_adj = ifelse(pval * length(mut_type) < 1, pval * length(mut_type), 1))
  return(list("output.kw" = output.kw, "plots" = plots))
}

factors <- c("country", "subsite")
pval_cutoff <- 0.05
removehypermut <- TRUE

# Define threshold for mutation types
threshold <- list(SBS = 100000, DBS = 6000, ID = 5000)

# Panels a and c
output.kw <- data.frame()
for (mut_type in c("SBS", "DBS", "ID")) {
  dat <- data
  if (removehypermut == T && mut_type %in% c("SBS", "DBS", "ID")) {
    dat <- dat %>% filter(get(paste0(mut_type, "_burden")) < threshold[[mut_type]])
  }
  for (f in factors) {
    output <- plot_kw_fct(dat, mut_type = paste0(mut_type, "_burden"), factors = f, pval_cutoff = 1, y_labs = "Mutation burden")

    output.kw <- rbind(output.kw, output$output.kw)

    ggsave(paste0("output/ExtendedDataFigure_1_", mut_type, "_", f, "_", Sys.Date(), ".pdf"), device = "pdf", width = 3.5, height = 4, dpi = 700)
  }
}

# Panel b - Faceted by subsites
factors <- "tobacco"
output.kw <- data.frame()

for (mut_type in c("SBS", "DBS", "ID")) {
  dat <- data
  if (removehypermut == T && mut_type %in% c("SBS", "DBS", "ID")) {
    dat <- dat %>% filter(get(paste0(mut_type, "_burden")) < threshold[[mut_type]])
  }
  for (f in factors) {
    output <- plot_kw_fct(dat, mut_type = paste0(mut_type, "_burden"), factors = f, pval_cutoff = 1, y_labs = "Mutation burden", facet_by = "subsite")

    output.kw <- rbind(output.kw, output$output.kw)

    ggsave(paste0("output/ExtendedDataFigure_1b_", mut_type, "_", f, "_", Sys.Date(), ".pdf"),
      device = "pdf",
      width = 7.2, height = 4.2, dpi = 700
    )
  }
}
